package ssh

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"my.com/mygo/utils/common"
)

const (
	hashMagic string = "|1|"
)

// Server SSH 登录信息
// 如果没有提供密码，则进行免密登录
type Server struct {
	Host     string
	Port     string
	User     string
	Password string
}

// Run run command
func (that *Server) Run(cmd string, args ...string) (string, error) {

	client, err := that.conn()
	if err != nil {
		return "", err
	}

	session, err := client.NewSession()
	if err != nil {
		return "", fmt.Errorf("new session error: %s", err.Error())
	}
	defer session.Close()
	defer client.Close()

	var b bytes.Buffer
	session.Stdout = &b
	commandLine := cmd
	for _, arg := range args {
		commandLine += " " + arg
	}
	err = session.Run(commandLine)
	if err != nil {
		return "", fmt.Errorf("Failed to run: %s", err.Error())
	}
	return b.String(), nil
}

func (that *Server) conn() (*ssh.Client, error) {
	// 建立 SSH 客户端连接 - 免密登录
	var configs []*ssh.ClientConfig
	var err error
	if that.Password == "" {
		configs, err = that.getConnPublicKeyConfig()
		if err != nil {
			return nil, err
		}
	} else {
		configs = append(configs, that.getConnPasswordConfig())
	}

	var client *ssh.Client
	for _, config := range configs {
		client, err = ssh.Dial("tcp", that.Host+":"+that.Port, config)
		if err == nil {
			break
		}
	}

	if err != nil {
		return nil, err
	}

	// 建立新会话
	return client, nil
}

func (that *Server) getConnPasswordConfig() *ssh.ClientConfig {
	return &ssh.ClientConfig{
		User: that.User,
		Auth: []ssh.AuthMethod{
			ssh.Password(that.Password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         10 * time.Second,
	}
}
func (that *Server) getHostKey() ([]*ssh.PublicKey, error) {
	// Every client must provide a host key check.  Here is a
	// simple-minded parse of OpenSSH's known_hosts file
	userHome, err := os.UserHomeDir()
	if err != nil {
		return nil, err
	}
	file, err := os.Open(filepath.Join(userHome, ".ssh", "known_hosts"))
	if err != nil {
		return nil, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	var hostKeys []*ssh.PublicKey
	for scanner.Scan() {
		fields := strings.Split(scanner.Text(), " ")
		if len(fields) != 3 {
			continue
		}

		if isRemoteHost(fields[0], that.Host) {
			var hostKey, _, _, _, err = ssh.ParseAuthorizedKey(scanner.Bytes())
			if err != nil {
				return nil, fmt.Errorf("error parsing %q: %v", fields[2], err)
			}
			hostKeys = append(hostKeys, &hostKey)
		}
	}

	if len(hostKeys) == 0 {
		return nil, fmt.Errorf("no hostkey for %s", that.Host)
	}

	return hostKeys, nil
}

func isRemoteHost(sshField string, host string) bool {
	if strings.HasPrefix(sshField, hashMagic) {
		items := strings.Split(sshField, "|")
		if len(items) != 4 {
			log.Println("ssh:hashed:length is not 4")
			return false
		}
		return common.SSHHash(items[2], host) == items[3]
	}
	return strings.Contains(sshField, host)
}

func getSigner() (*ssh.Signer, error) {

	userHome, err := os.UserHomeDir()
	if err != nil {
		log.Panic(err)
	}

	// A public key may be used to authenticate against the remote
	// server by using an un-encrypted PEM-encoded private key file.
	//
	// If you have an encrypted private key, the crypto/x509 package
	// can be used to decrypt it.
	key, err := os.ReadFile(filepath.Join(userHome, ".ssh", "id_rsa"))
	if err != nil {
		return nil, fmt.Errorf("unable to read private key: %v", err)
	}

	// Create the Signer for this private key.
	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		return nil, fmt.Errorf("unable to parse private key: %v", err)
	}

	return &signer, nil
}

func (that *Server) getConnPublicKeyConfig() ([]*ssh.ClientConfig, error) {

	// Every client must provide a host key check.  Here is a
	// simple-minded parse of OpenSSH's known_hosts file
	var configs []*ssh.ClientConfig
	hostKeys, err := that.getHostKey()
	if err != nil {
		return nil, err
	}

	signer, err := getSigner()
	if err != nil {
		return nil, err
	}

	for _, hostKey := range hostKeys {
		configs = append(configs, &ssh.ClientConfig{
			User: that.User,
			Auth: []ssh.AuthMethod{
				// Use the PublicKeys method for remote authentication.
				ssh.PublicKeys(*signer),
			},
			HostKeyCallback: ssh.FixedHostKey(*hostKey),
			Timeout:         10 * time.Second,
		})
	}

	return configs, nil
}
